{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "340dcc54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "from plyfile import PlyData, PlyElement\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "518dd393",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'C:\\\\Users\\\\amira\\\\Documents\\\\1 NJ Transit Project\\\\3 PC Labeling\\\\4 PTS and Labels'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get current directory\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f6eacb63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect with directories\n",
    "ID_dir = \"C:/Users/amira/Documents/1 NJ Transit Project/3 PC Labeling/3 Component Extraction/\"\n",
    "points_save_dir = \"C:/Users/amira/Documents/1 NJ Transit Project/3 PC Labeling/4 PTS and Labels/Points\"\n",
    "labels_save_dir = \"C:/Users/amira/Documents/1 NJ Transit Project/3 PC Labeling/4 PTS and Labels/Labels\"\n",
    "FILE = os.listdir(ID_dir)\n",
    "Categories = [\"Structure\",\"Ties\",\"Tracks\",\"Others\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fc5bd86e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analyzing dataset #0...\n",
      "Analyzing dataset #1...\n",
      "Analyzing dataset #2...\n",
      "Analyzing dataset #3...\n",
      "Analyzing dataset #4...\n",
      "Analyzing dataset #5...\n",
      "Analyzing dataset #6...\n",
      "Analyzing dataset #7...\n",
      "Analyzing dataset #8...\n",
      "Analyzing dataset #9...\n",
      "Analyzing dataset #10...\n",
      "Analyzing dataset #11...\n",
      "Analyzing dataset #12...\n",
      "Analyzing dataset #13...\n",
      "Analyzing dataset #14...\n",
      "Analyzing dataset #15...\n",
      "Analyzing dataset #16...\n",
      "Analyzing dataset #17...\n",
      "Analyzing dataset #18...\n",
      "Analyzing dataset #19...\n",
      "Analyzing dataset #20...\n",
      "Analyzing dataset #21...\n",
      "Analyzing dataset #22...\n",
      "Analyzing dataset #23...\n",
      "Analyzing dataset #24...\n",
      "Analyzing dataset #25...\n",
      "Analyzing dataset #26...\n",
      "Analyzing dataset #27...\n",
      "Analyzing dataset #28...\n",
      "Analyzing dataset #29...\n",
      "Analyzing dataset #30...\n",
      "Analyzing dataset #31...\n",
      "Analyzing dataset #32...\n",
      "Analyzing dataset #33...\n",
      "Analyzing dataset #34...\n",
      "Analyzing dataset #35...\n",
      "Analyzing dataset #36...\n",
      "Analyzing dataset #37...\n",
      "Analyzing dataset #38...\n",
      "Analyzing dataset #39...\n"
     ]
    }
   ],
   "source": [
    "#%%%%%%%%%%%%%%%%%%% Processing data %%%%%%%%%%%%%%%%%%%%%%\n",
    "for j in range(40):\n",
    "    print(\"Analyzing dataset #\"+str(j)+\"...\")\n",
    "    os.chdir(ID_dir+FILE[j])\n",
    "    if os.path.exists(Categories[0]+'.ply'):\n",
    "        with open(Categories[0]+'.ply', 'rb') as f:\n",
    "            open_data = PlyData.read(f)\n",
    "            Structure_data = np.concatenate((open_data['vertex']['x'].reshape(-1,1),open_data['vertex']['y'].reshape(-1,1),open_data['vertex']['z'].reshape(-1,1)),axis = 1) \n",
    "            Structure_label = np.zeros((Structure_data.shape[0],4))\n",
    "            Structure_label[:,0:1] = np.ones((Structure_data.shape[0],1))\n",
    "\n",
    "    if os.path.exists(Categories[1]+'.ply'):\n",
    "        with open(Categories[1]+'.ply', 'rb') as f:\n",
    "            open_data = PlyData.read(f)\n",
    "            Ties_data = np.concatenate((open_data['vertex']['x'].reshape(-1,1),open_data['vertex']['y'].reshape(-1,1),open_data['vertex']['z'].reshape(-1,1)),axis = 1)\n",
    "            Ties_label = np.zeros((Ties_data.shape[0],4))\n",
    "            Ties_label[:,1:2] = np.ones((Ties_data.shape[0],1))\n",
    "\n",
    "    if os.path.exists(Categories[2]+'.ply'):\n",
    "        with open(Categories[2]+'.ply', 'rb') as f:\n",
    "            open_data = PlyData.read(f)\n",
    "            Tracks_data = np.concatenate((open_data['vertex']['x'].reshape(-1,1),open_data['vertex']['y'].reshape(-1,1),open_data['vertex']['z'].reshape(-1,1)),axis = 1)\n",
    "            Tracks_label = np.zeros((Tracks_data.shape[0],4))\n",
    "            Tracks_label[:,2:3] = np.ones((Tracks_label.shape[0],1))\n",
    "\n",
    "    if os.path.exists(Categories[3]+'.ply'):    \n",
    "        with open(Categories[3]+'.ply', 'rb') as f:\n",
    "            open_data = PlyData.read(f)\n",
    "            Others_data = np.concatenate((open_data['vertex']['x'].reshape(-1,1),open_data['vertex']['y'].reshape(-1,1),open_data['vertex']['z'].reshape(-1,1)),axis = 1)\n",
    "            Others_label = np.zeros((Others_data.shape[0],4))\n",
    "            Others_label[:,3:] = np.ones((Others_data.shape[0],1))\n",
    "\n",
    "    if not os.path.exists(Categories[2]+'.ply'):    \n",
    "        Data = np.concatenate((Structure_data, Ties_data, Others_data), axis = 0)\n",
    "        Label = np.concatenate((Structure_label, Ties_label, Others_label), axis = 0)\n",
    "    if not os.path.exists(Categories[3]+'.ply'):    \n",
    "        Data = np.concatenate((Structure_data, Ties_data, Tracks_data), axis = 0)\n",
    "        Label = np.concatenate((Structure_label, Ties_label, Tracks_label), axis = 0)    \n",
    "    else:     \n",
    "        Data = np.concatenate((Structure_data, Ties_data, Tracks_data, Others_data), axis = 0)\n",
    "        Label = np.concatenate((Structure_label, Ties_label, Tracks_label, Others_label), axis = 0)\n",
    "    LabeledData = np.concatenate((Data, Label), axis = 1)\n",
    "    len = LabeledData.shape[0]\n",
    "    index = np.arange(0,len)\n",
    "    np.random.shuffle(index)\n",
    "    LabeledData_shuffled = np.zeros((LabeledData.shape[0],LabeledData.shape[1]))\n",
    "    for i in range(len):\n",
    "        LabeledData_shuffled[i,:] = LabeledData[index[i],:]\n",
    "\n",
    "    # Colors\n",
    "#     Colors = {}\n",
    "#     for i in range(len):\n",
    "#         if LabeledData_shuffled [i,3] == 1:\n",
    "#             Colors[i] = '#b41f1f'\n",
    "#         elif LabeledData_shuffled [i,4] == 1:\n",
    "#             Colors[i] = '#1fb421'\n",
    "#         elif LabeledData_shuffled [i,5] == 1:\n",
    "#             Colors[i] = '#e8e113'\n",
    "#         else:\n",
    "#             Colors[i] = '#1f4cb4'\n",
    "    os.chdir(\"C:/Users/amira/\\Documents/1 NJ Transit Project/3 PC Labeling/4 PTS and Labels/Figures\")\n",
    "\n",
    "    #%%%%%%%%%%%%%%%%% Visualize data %%%%%%%%%%%%%%%%%%%\n",
    "#     fig = plt.figure(figsize=(12, 6))\n",
    "#     ax = fig.gca(projection='3d') \n",
    "#     for i in range(1,len,200):\n",
    "#         ax.scatter(LabeledData_shuffled[i,0], LabeledData_shuffled[i,2], LabeledData_shuffled[i,1], c=Colors[i])\n",
    "#     plt.ylim(-4000,1000)\n",
    "#     plt.show()\n",
    "#     fig.savefig('Fig_'+FILE[j]+'.png')\n",
    "\n",
    "    #%%%%%%%%%%%%%%%%% Saving data %%%%%%%%%%%%%%%%%%%%%%%\n",
    "    os.chdir(points_save_dir)    \n",
    "    Data_unst = LabeledData_shuffled[:,0:3]\n",
    "    a1_file = open(\"data_\"+FILE[j]+\".PTS\", \"w\")\n",
    "    np.savetxt(a1_file, Data_unst)\n",
    "    a1_file.close()\n",
    "\n",
    "    # Saving labels\n",
    "    Label_unst = LabeledData_shuffled[:,3:]\n",
    "    os.chdir(labels_save_dir+\"/\"+Categories[0])  \n",
    "    Structure_label_unst = Label_unst[:,0:1]\n",
    "    a2_file = open(\"label_\"+FILE[j]+\".SEG\", \"w\")\n",
    "    np.savetxt(a2_file, Structure_label_unst)\n",
    "    a2_file.close()\n",
    "\n",
    "    os.chdir(labels_save_dir+\"/\"+Categories[1]) \n",
    "    Ties_label_unst = Label_unst[:,1:2]\n",
    "    a3_file = open(\"label_\"+FILE[j]+\".SEG\", \"w\")\n",
    "    np.savetxt(a3_file, Ties_label_unst)\n",
    "    a3_file.close()\n",
    "\n",
    "    os.chdir(labels_save_dir+\"/\"+Categories[2]) \n",
    "    Tracks_label_unst = Label_unst[:,2:3]\n",
    "    a4_file = open(\"label_\"+FILE[j]+\".SEG\", \"w\")\n",
    "    np.savetxt(a4_file, Tracks_label_unst)\n",
    "    a4_file.close()\n",
    "\n",
    "    os.chdir(labels_save_dir+\"/\"+Categories[3]) \n",
    "    Others_label_unst = Label_unst[:,3:4]\n",
    "    a5_file = open(\"label_\"+FILE[j]+\".SEG\", \"w\")\n",
    "    np.savetxt(a5_file, Others_label_unst)\n",
    "    a5_file.close()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
